import os
import torch
import subprocess
from pprint import pprint

from .zenodo import Zenodo

PLEASE_PUBLISH = (
    "\nDon't forget to share your pretrained models at "
    "https://zenodo.org/communities/asteroid-models/ ! =)\n"
    "You can directly use our CLI for that, run this: \n"
    "`asteroid-upload {} --uploader \"Your name here\"`\n"
)

HREF = '<a href="{}">{}</a>'
CC_SA = 'Attribution-ShareAlike 3.0 Unported'
CC_SA_LINK = 'https://creativecommons.org/licenses/by-sa/3.0/'
ASTEROID_REF = HREF.format('https://github.com/mpariente/asteroid', 'Asteroid')


def save_publishable(publish_dir, model_dict, metrics=None, train_conf=None):
    """ Save models to prepare for publication / model sharing.

    Args:
        publish_dir (str): Path to the publishing directory.
            Usually under exp/exp_name/publish_dir
        model_dict (dict): dict at least with keys `model_args`,
            `state_dict`,`dataset` or `licenses`
        metrics (dict): dict with evaluation metrics.
        train_conf (dict): Training configuration dict (from conf.yml).

    Returns:
        dict, same as `model_dict` with added fields.

    Raises:
        AssertionError when either `model_args`, `state_dict`,`dataset` or
            `licenses` are not present is `model_dict.keys()`
    """
    assert 'model_args' in model_dict.keys(), "`model_args` not found in model dict."
    assert 'state_dict' in model_dict.keys(), "`state_dict` not found in model dict."
    assert 'dataset' in model_dict.keys(), "`dataset` not found in model dict."
    assert 'licenses' in model_dict.keys(), "`licenses` not found in model dict."
    assert isinstance(metrics, dict), "Cannot upload a model without metrics."
    # Additional infos.
    if os.path.exists(os.path.join(publish_dir, 'recipe_name.txt')):
        recipe_name = next(open(os.path.join(publish_dir, 'recipe_name.txt')))
        recipe_name.replace('\n', '')  # remove next line
    else:
        recipe_name = 'Unknown'
    model_dict['infos']['recipe_name'] = recipe_name
    model_dict['infos']['training_config'] = train_conf
    model_dict['infos']['final_metrics'] = metrics
    torch.save(model_dict, os.path.join(publish_dir, 'model.pth'))
    print(PLEASE_PUBLISH.format(publish_dir))
    return model_dict


def upload_publishable(publish_dir, uploader=None, affiliation=None,
                       git_username=None, token=None, force_publish=False,
                       use_sandbox=False, unit_test=False):
    """ Entry point to upload publishable model.

    Args:
        publish_dir (str): Path to the publishing directory.
            Usually under exp/exp_name/publish_dir
        uploader (str): Full name of the uploader (Ex: Manuel Pariente)
        affiliation (str, optional): Affiliation (no accent).
        git_username (str, optional): GitHub username.
        token (str): Access token generated to upload depositions.
        force_publish (bool): Whether to directly publish without
            asking confirmation before. Defaults to False.
        use_sandbox (bool): Whether to use Zenodo's sandbox instead of
            the official Zenodo.
        unit_test (bool): If True, we do not ask user input and do not publish.

    """
    def get_answer():
        out = input('\n\nDo you want to publish it now (irreversible)? y/n'
                    '(Recommended: n).\n')
        if out not in ['y', 'n']:
            print(f'\nExpected one of [`y`, `n`], received {out}, please retry.')
            return get_answer()
        return out

    if uploader is None:
        raise ValueError('Need uploader name')

    # Make publishable model and save it
    model_path = os.path.join(publish_dir, 'model.pth')
    publish_model_path = os.path.join(publish_dir, 'published_model.pth')
    model = torch.load(model_path)
    model = _populate_publishable(
        model,
        uploader=uploader,
        affiliation=affiliation,
        git_username=git_username,
    )
    torch.save(model, publish_model_path)

    # Get Zenodo access token
    if token is None:
        token = os.getenv('ACCESS_TOKEN')
        if token is None:
            raise ValueError(
                'Need an access token to Zenodo to upload the model. Either '
                'set ACCESS_TOKEN environment variable or pass it directly '
                '(`asteroid-upload --token ...`).'
                'If you do not have a access token, first create a Zenodo '
                'account (https://zenodo.org/signup/), create a token '
                'https://zenodo.org/account/settings/applications/tokens/new/'
                'and you are all set to help us ! =)'
            )

    # Do the actual upload
    zen, dep_id = zenodo_upload(model, token, model_path=publish_model_path,
                                use_sandbox=use_sandbox)
    address = os.path.join(zen.zenodo_address, 'deposit', str(dep_id))
    if force_publish:
        r_publish = zen.publish_deposition(dep_id)
        pprint(r_publish.json())
        print("You can also visit it at {}".format(address))
        return r_publish
    # Give choice
    current = zen.get_deposition(dep_id)
    print(f"\n\n This is the current state of the deposition "
          f"(see here {address}): ")
    pprint(current.json())
    # Patch to run unit test
    if unit_test:
        return zen, current
    else:
        inp = get_answer()
    # Get user input
    if inp == 'y':
        _ = zen.publish_deposition(dep_id)
        print("Visit it at {}".format(address))
    else:
        print(f'Did not finalize the upload, please visit {address} to finalize '
              f'it.')


def _populate_publishable(model, uploader=None, affiliation=None,
                          git_username=None):
    """ Populate infos in publishable model.

    Args:
        model (dict): Model to publish, with `infos` key, at least.
        uploader (str): Full name of the uploader (Ex: Manuel Pariente)
        affiliation (str, optional): Affiliation (no accent).
        git_username (str, optional): GitHub username.

    Returns:
        dict (model), same as input `model`

    Notes:
        If a `git_username` is not specified, we look for it somehow, or take
        the laptop username.
    """
    # Get username somehow
    if git_username is None:
        git_username = get_username()

    # Example: mpariente/ConvTasNet_WHAM_sepclean
    model_name = '_'.join([model['model_name'], model['dataset'],
                           model['task'].replace('_', '')])
    upload_name = git_username + '/' + model_name
    # Write License Notice
    license_note = make_license_notice(model_name, model['licenses'],
                                       uploader=uploader)
    # Add infos
    model['infos']['uploader'] = uploader
    model['infos']['git_username'] = git_username
    model['infos']['affiliation'] = affiliation if affiliation else "Unknown"
    model['infos']['upload_name'] = upload_name
    model['infos']['license_note'] = license_note
    return model


def get_username():
    """ Get git of FS username for upload. """
    username = subprocess.check_output(["git", "config", "user.name"])
    username = username.decode('utf-8')[:-1]
    if not username:  # Empty string
        import getpass
        username = getpass.getuser()
    return username


def make_license_notice(model_name, licenses, uploader=None):
    """ Make license notice based on license dicts.

    Args:
        model_name (str): Name of the model.
        licenses (List[dict]): List of dict with
            keys (`title`, `title_link`, `author`, `author_link`,
                  `licence`, `licence_link`).
        uploader (str): Name of the uploader such as "Manuel Pariente".

    Returns:
        str, the license note describing the model, it's attribution,
            the original licenses, what we license it under and the licensor.
    """
    if uploader is None:
        raise ValueError("Cannot share model without uploader.")
    note = "This work \"{}\" is a derivative ".format(model_name)
    for l_dict in licenses:
        # Clickable links in HTML.
        title = HREF.format(l_dict['title_link'], l_dict['title'])
        author = HREF.format(l_dict['author_link'], l_dict['author'])
        license_h = HREF.format(l_dict['license_link'], l_dict['license'])
        comm = " (Research only)" if l_dict['non_commercial'] else ""
        note += f"of {title} by {author}, used under {license_h}{comm}"
        note += "; "
    note = note[:-2] + '. '  # Remove the last ;
    cc_sa = HREF.format(CC_SA_LINK, CC_SA)
    note += f"\"{model_name}\" is licensed under {cc_sa} by {uploader}."
    return note


def zenodo_upload(model, token, model_path=None, use_sandbox=False):
    """ Create deposit and upload metadata + model

    Args:
        model (dict):
        token (str): Access token.
        model_path (str): Saved model path.
        use_sandbox (bool): Whether to use Zenodo's sandbox instead of
            the official Zenodo.

    Returns:
        Zenodo (Zenodo instance with access token)
        int (deposit ID)

    Notes:
        If `model_path` is not specified, save the model in tmp.pth and
        remove it after upload.
    """
    model_path_was_none = False
    if model_path is None:
        model_path_was_none = True
        model_path = 'tmp.pth'
        torch.save(model, model_path)
        # raise ValueError("Need path")

    zen = Zenodo(token, use_sandbox=use_sandbox)
    metadata = make_metadata_from_model(model)
    r = zen.create_new_deposition(metadata=metadata)
    if r.status_code != 200:
        print(r.json())
        raise RuntimeError('Could not create the deposition, check the '
                           'provided token.')
    dep_id = r.json()["id"]
    _ = zen.upload_new_file_to_deposition(dep_id, model_path, name='model.pth')
    if model_path_was_none:
        os.remove(model_path)
    return zen, dep_id


def make_metadata_from_model(model):
    """ Create Zenodo deposit metadata for a given publishable model.
    Args:
        model (dict): Dictionary with all infos needed to publish.
            More info to come.

    Returns:
        dict, the metadata to create the Zenodo deposit with.

    Notes:
        We remove the PESQ from the final results as a license is needed to
        use it.
    """
    infos = model['infos']
    # Description section
    description = '<p><strong>Description: </strong></p>'
    tmp = 'This model was trained by {} using the {} recipe in {}. '
    description += tmp.format(infos['uploader'], infos['recipe_name'],
                              ASTEROID_REF)
    tmp = '</a>It was trained on the <code>{}</code> task of the {} dataset.</p>'
    description += tmp.format(model['task'], model['dataset'])

    # Training config section
    description += '<p>&nbsp;</p>'
    description += '<p><strong>Training config:</strong></p>'
    description += two_level_dict_html(infos['training_config'])

    # Results section
    description += '<p>&nbsp;</p>'
    description += '<p><strong>Results:</strong></p>'
    display_result = {
        k: v for k, v in infos['final_metrics'].items()
        if 'pesq' not in k.lower()
    }
    description += display_one_level_dict(display_result)

    # License section
    description += '<p>&nbsp;</p>'
    description += '<p><strong>License notice:</strong></p>'
    description += infos['license_note']

    # Putting it together.
    metadata = {
        'title': infos['upload_name'],
        'upload_type': 'software',
        'description': description,
        'creators': [{'name': infos['uploader'],
                      'affiliation': infos['affiliation']}],
        'communities': [{'identifier': 'zenodo'},
                        {'identifier': 'asteroid-models'}],
        'keywords': ['Asteroid', 'audio source separation', model['dataset'],
                     model['task'], model['model_name'], 'pretrained model'],
        'license': 'CC-BY-SA-3.0'
    }
    return metadata


def two_level_dict_html(dic):
    """ Two-level dict to HTML.
    Args:
        dic (dict): two-level dict

    Returns:
        str for HTML-encoded two level dic
    """
    html = "<ul>"
    for k in dic.keys():
        # Open field
        html += f'<li>{k}: <ul>'
        for k2 in dic[k].keys():
            val = str(dic[k][k2])
            html += f'<li>{k2}: {val}</li>'
        # Close field
        html += '</il></ul>'
    html += '</ul>'
    return html


def display_one_level_dict(dic):
    """ Single level dict to HTML
    Args:
        dic (dict):

    Returns:
        str for HTML-encoded single level dic
    """
    html = "<ul>"
    for k in dic.keys():
        # Open field
        val = str(dic[k])
        html += f'<li>{k}: {val} </li>'
    html += '</ul>'
    return html
